# ========================================================================================================================================
# Copyright (c) 2025 CompanyName. All rights reserved.
# Author:         22020873 陈泽欣
# Project:        Design of Deep Learning Fundamental Course
# Module:         posenet_detect.py
# Date:           2025-05-24
# Description:    本模块定义并实现了 PoseNet 网络模型，用于骰子姿态估计中的旋转角度预测。
#                 包含以下核心功能：
#                 - 基于 CNN 和残差结构的轻量级特征提取网络；
#                 - 回归输出旋转向量（rvec）用于姿态估计；
#                 - 图像预处理与推理接口封装，便于集成到整体系统中；
#                 是整个骰子姿态识别系统中深度学习预测部分的核心组件。
# ========================================================================================================================================

import torch
import torch.nn as nn
from torchvision import transforms
import cv2
import random

# ==========================================================
#                        残差模块（ResBlock）
# ==========================================================

class ResBlock(nn.Module):
    def __init__(self, in_channels):
        super(ResBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(in_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(in_channels)

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        out += identity  # 跳跃连接
        out = self.relu(out)

        return out

# ==========================================================
#                        模型定义 PoseNet
# ==========================================================
class PoseNet(nn.Module):
    def __init__(self):
        super(PoseNet, self).__init__()
        # Backbone: VGG-style CNN
        self.features = nn.Sequential(
            # Block 1
            nn.Conv2d(1, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            ResBlock(64),
            nn.MaxPool2d(kernel_size=2, stride=2),

            # Block 2
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),

            # Block 3
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            ResBlock(256),
            nn.MaxPool2d(kernel_size=2, stride=2),

            # Block 4
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            # ResBlock(512),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )

        self.regressor = nn.Sequential(
            nn.Flatten(),
            nn.Linear(512 * 8 * 8, 4096),
            nn.ReLU(True),
            nn.Dropout(p=0.5),
            nn.Linear(4096, 2048),
            nn.ReLU(True),
            nn.Dropout(p=0.5),
            nn.Linear(2048, 3)  # 直接输出原始 rvec（弧度）
        )

    def forward(self, x):
        x = self.features(x)
        rvec = self.regressor(x)
        return rvec


# ==========================================================
#                        推理函数
# ==========================================================
def predict_rvec(model, image, device='cuda'):
    # 预处理
    transform = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((128, 128)),
        transforms.Grayscale(num_output_channels=1),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5], std=[0.5]),
    ])

    image_tensor = transform(image).unsqueeze(0).to(device)

    # 推理
    model.eval()
    with torch.no_grad():
        output = model(image_tensor)

    rvec_pred = output.squeeze().cpu().numpy()

    return rvec_pred

